#!/bin/bash
# 磁盘自动化分区、格式化及挂载脚本（增强版V2）
# 功能：支持命令行参数指定磁盘设备和挂载点，自动识别NVMe设备，增加安全卸载逻辑
# 日期：2025-08-08

# 基本用法（使用默认挂载点）
#./disk.sh /dev/nvme0n1
# 指定自定义挂载点
#./disk.sh /dev/sdb /mnt/archive
# 查看可用磁盘（无参数时）
#./disk.sh

# 检查root权限
if [ "$(id -u)" != "0" ]; then
    echo "错误：必须使用root权限运行此脚本！" >&2
    exit 1
fi

# 使用命令行参数指定磁盘设备和挂载点
if [ $# -lt 1 ]; then
    echo "用法: $0 <磁盘设备路径> [挂载点]"
    echo "示例: $0 /dev/sdb"
    echo "       $0 /dev/nvme0n1 /mnt/data2"
    echo "可用磁盘列表："
    lsblk -d -o NAME,SIZE,MODEL | grep -v 'loop' | sed 's/^/  /'
    exit 1
fi

DISK_DEVICE="$1"
MOUNT_POINT="${2:-/data}"  # 默认挂载点为/data[6](@ref)

# 验证设备存在性及类型
validate_disk() {
    # 检查设备是否存在
    if [ ! -e "$DISK_DEVICE" ]; then
        echo "错误：设备 $DISK_DEVICE 不存在！" >&2
        echo "可用磁盘列表：" >&2
        lsblk -d -o NAME,SIZE,MODEL | grep -v 'loop' | sed 's/^/  /' >&2
        exit 1
    fi

    # 使用lsblk验证设备类型（磁盘/分区）
    if ! lsblk -d "$DISK_DEVICE" &>/dev/null; then
        echo "错误：$DISK_DEVICE 不是磁盘设备（可能是分区）！请指定整个磁盘（如 /dev/sdb 或 /dev/nvme0n1）。" >&2
        exit 1
    fi
}
validate_disk

# 安全卸载函数
safe_unmount() {
    # 卸载所有关联挂载点
    grep "$DISK_DEVICE" /proc/mounts | awk '{print $2}' | while read -r mount_point; do
        echo "卸载挂载点: $mount_point"
        umount -f "$mount_point" 2>/dev/null
    done

    # 关闭关联的swap分区
    swapon --show=NAME --noheadings | grep "^$DISK_DEVICE" | while read -r swap_part; do
        echo "关闭swap分区: $swap_part"
        swapoff "$swap_part" 2>/dev/null
    done

    # 移除内核中的分区表缓存
    partx -d "$DISK_DEVICE" 2>/dev/null
}

# 检查挂载点状态
check_mount_point() {
    # 检查挂载点是否已被使用
    if mount | grep -q " on $MOUNT_POINT "; then
        echo "错误：挂载点 $MOUNT_POINT 已被使用！" >&2
        exit 1
    fi

    # 检查挂载点目录状态
    if [ -d "$MOUNT_POINT" ]; then
        if [ "$(ls -A $MOUNT_POINT)" ]; then
            echo "错误：挂载点 $MOUNT_POINT 非空目录！" >&2
            exit 1
        fi
    else
        echo "创建挂载点目录: $MOUNT_POINT"
        mkdir -p "$MOUNT_POINT" || {
            echo "错误：无法创建挂载点目录！" >&2
            exit 1
        }
    fi
}

# 获取磁盘大小（字节）
DISK_SIZE=$(blockdev --getsize64 "$DISK_DEVICE")
DISK_SIZE_GB=$((DISK_SIZE / 1024**3))  # 转换为GB

# 安全卸载现有分区
safe_unmount

# 清除磁盘签名
echo "清除磁盘签名..."
wipefs -a -f "$DISK_DEVICE" >/dev/null
dd if=/dev/zero of="$DISK_DEVICE" bs=1M count=100 status=none >/dev/null
sync

# 选择分区方案
if [ "$DISK_SIZE_GB" -lt 2000 ]; then
    echo "检测到磁盘大小 ${DISK_SIZE_GB}GB < 2TB，使用 MBR(fdisk) 分区"
    echo -e "o\nn\np\n1\n\n\nw" | fdisk "$DISK_DEVICE" >/dev/null 2>&1
else
    echo "检测到磁盘大小 ${DISK_SIZE_GB}GB ≥ 2TB，使用 GPT(parted) 分区"
    parted -s "$DISK_DEVICE" mklabel gpt >/dev/null
    parted -s "$DISK_DEVICE" mkpart primary 0% 100% >/dev/null
fi

# 确定分区名称（处理NVMe设备）
if [[ "$DISK_DEVICE" =~ /dev/nvme[0-9]+n[0-9]+$ ]]; then
    PARTITION="${DISK_DEVICE}p1"
else
    PARTITION="${DISK_DEVICE}1"
fi

# 等待分区识别
echo "等待分区设备初始化..."
count=0
until [ -e "$PARTITION" ] || [ $count -ge 10 ]; do
    sleep 1
    partprobe "$DISK_DEVICE" >/dev/null 2>&1
    ((count++))
done

if [ ! -e "$PARTITION" ]; then
    echo "错误：分区 $PARTITION 未创建！" >&2
    exit 1
fi

# 格式化分区为XFS
echo "正在格式化分区 $PARTITION 为 XFS 文件系统..."
mkfs.xfs -f -L "DATA_DISK" "$PARTITION" >/dev/null 2>&1
if [ $? -ne 0 ]; then
    echo "格式化失败！请检查磁盘状态。" >&2
    exit 1
fi

# 检查挂载点状态
check_mount_point

# 获取分区UUID
UUID=$(blkid -s UUID -o value "$PARTITION")

# 永久挂载配置（安全更新方式）
sed -i "\|^[^#].*$MOUNT_POINT |d" /etc/fstab
echo "UUID=$UUID $MOUNT_POINT xfs defaults,noatime,nodiratime 0 0" >> /etc/fstab

# 挂载分区
mount -a >/dev/null 2>&1

# 验证结果
if mount | grep -q "$MOUNT_POINT"; then
    echo "成功：分区已挂载到 $MOUNT_POINT"
    df -hT "$MOUNT_POINT"
    echo "磁盘配置已更新至 /etc/fstab"
    echo "挂载选项：noatime,nodiratime（已优化I/O性能）"
else
    echo "挂载失败！请检查 /etc/fstab 配置" >&2
    # 回滚fstab修改
    sed -i "\|^[^#].*$MOUNT_POINT |d" /etc/fstab
    exit 1
fi
